import pandas as pd

import input

# 数据预处理和特征工程
def preprocess_data(dataframes):
    medal_counts = dataframes["medal_counts"]
    hosts = dataframes["hosts"]
    programs = dataframes["programs"]
    athletes = dataframes["athletes"]

    # 数据预处理， 第一步删除
    medal_counts = medal_counts[medal_counts['NOC'] != 'Mixed team']
    medal_counts = medal_counts[medal_counts['NOC'] != 'Chinese Taipei']
    medal_counts = medal_counts[medal_counts['NOC'] != 'Formosa']
    medal_counts = medal_counts[medal_counts['NOC'] != 'Yugoslavia']
    medal_counts = medal_counts[~((medal_counts['NOC'] == 'United Team of Germany') & (medal_counts['Year'] == 1960))]
    medal_counts = medal_counts[~((medal_counts['NOC'] == 'Hong Kong') & (medal_counts['Year'] < 2020))]
    medal_counts = medal_counts[~((medal_counts['NOC'] == 'Ivory Coast') & (medal_counts['Year'] == 2024))]
    medal_counts = medal_counts[~((medal_counts['NOC'] == 'Turkey') & (medal_counts['Year'] == 2024))]
    # 筛选出NOC为"Hong Kong"且Year大于2019的数据
    hong_kong_data = medal_counts[(medal_counts['NOC'] == 'Hong Kong') & (medal_counts['Year'] > 2019)]
    # 筛选出NOC为"China"的数据
    china_data = medal_counts[medal_counts['NOC'] == 'China']
    # 遍历筛选出的Hong Kong数据
    for index, row in hong_kong_data.iterrows():
        year = row['Year']
        # 找到对应的China数据
        china_row = china_data[china_data['Year'] == year]
        if not china_row.empty:
            # 将Hong Kong的奖牌数量加到China的奖牌数量中
            medal_counts.loc[china_row.index, 'Gold'] += row['Gold']
            medal_counts.loc[china_row.index, 'Silver'] += row['Silver']
            medal_counts.loc[china_row.index, 'Bronze'] += row['Bronze']
            medal_counts.loc[china_row.index, 'Total'] += row['Total']
        else:
            # 如果没有对应的China数据，则直接将Hong Kong的数据改为China
            medal_counts.loc[index, 'NOC'] = 'China'
    # 删除原始的Hong Kong数据（已经合并的数据）
    medal_counts = medal_counts[~((medal_counts['NOC'] == 'Hong Kong') & (medal_counts['Year'] > 2019))]

    team_mapping = input.load_json("./json/team_mapping.json") # 使用自定义函数读取json文件
    athletes['Team'] = athletes['Team'].replace(team_mapping)
    noc_mapping = input.load_json("./json/noc_mapping.json") # 使用自定义函数读取json文件, Mixed team合并队伍
    medal_counts['NOC'] = medal_counts['NOC'].replace(noc_mapping)

    medal_counts = medal_counts.rename(columns={"NOC": "Country"})
    medal_counts["Country"] = medal_counts["Country"].str.replace(r"[^\x00-\x7F]+", "", regex=True).str.strip()
    hosts = hosts.rename(columns={"ï»¿Year": "Year", "Host": "Host_City"})
    hosts["Host_Country"] = hosts["Host_City"].str.extract(r",\s*([^\(]+)\s*(?:\(|$)", expand=False).str.strip()


    ##################################### 是否为主办国 ########################################
    medal_counts = medal_counts.merge(hosts[["Year", "Host_Country"]], on="Year", how="left")
    medal_counts["Host_Country"] = medal_counts["Host_Country"].str.replace(r"[^\x00-\x7F]+", "", regex=True)
    medal_counts["Is_Host"] = (medal_counts["Country"] == medal_counts["Host_Country"]).astype(int)


    ##################################### 上一次奖牌总数 ########################################
    medal_counts["Last_Total_Medals"] = medal_counts.groupby("Country")["Total"].shift(1).fillna(0)


    ##################################### 比赛项目数量 ##########################################
    # 清理列名中的特殊字符
    programs.columns = programs.columns.str.replace(r"[^a-zA-Z0-9\s]", "", regex=True)
    # 将年份列转换为整数类型
    year_columns = [col for col in programs.columns if col.isdigit()]
    programs[year_columns] = programs[year_columns].apply(pd.to_numeric, errors='coerce')
    # 提取最后三行数据
    all_rows = programs.iloc[-74:]
    # 将非数字字符或空格替换为0，并转换为数值类型
    for col in all_rows.columns:
        if col.isdigit():  # 检查列名是否为年份（数字）
            all_rows[col] = all_rows[col].replace(r'[^\d]+', 0, regex=True)  # 替换非数字字符为0
            all_rows[col] = pd.to_numeric(all_rows[col], errors='coerce').fillna(0)  # 转换为数值类型，非数字的填充为0
    # 提取年份列名
    year_columns = [col for col in programs.columns if col.isdigit()]
    # 初始化一个空的列表来存储每年的加权和
    weighted_sums = []
    # 初始化一个空的列表来存储每年每个运动分类的项目数量比例
    aquatics_sums = []
    ball_sports_sums = []
    equestrian_sums = []
    martial_arts_sums = []
    extreme_sports_sums = []
    other_sports_sums = []
    # 遍历每一年，计算加权和
    for year in year_columns:
        total_events = all_rows.iloc[71][year]

        total_disciplines = all_rows.iloc[72][year]
        total_sports = all_rows.iloc[73][year]

        aquatics = (all_rows.iloc[0][year] + all_rows.iloc[1][year] + all_rows.iloc[2][year] + all_rows.iloc[3][year] + all_rows.iloc[4][year] + all_rows.iloc[65][year] + all_rows.iloc[69][year] + all_rows.iloc[70][year] + all_rows.iloc[49][year] + all_rows.iloc[48][year] + all_rows.iloc[52][year])
        ball_sports =(all_rows.iloc[7][year] + all_rows.iloc[8][year] + all_rows.iloc[9][year] + all_rows.iloc[10][year] + all_rows.iloc[11][year] + all_rows.iloc[12][year] + all_rows.iloc[30][year] + all_rows.iloc[31][year] + all_rows.iloc[32][year] + all_rows.iloc[33][year] + all_rows.iloc[39][year] + all_rows.iloc[45][year] + all_rows.iloc[46][year] + all_rows.iloc[47][year] + all_rows.iloc[18][year] + all_rows.iloc[17][year] + all_rows.iloc[42][year] + all_rows.iloc[50][year] + all_rows.iloc[51][year] + all_rows.iloc[56][year] + all_rows.iloc[58][year] + all_rows.iloc[60][year] + all_rows.iloc[37][year] + all_rows.iloc[38][year])
        equestrian = (all_rows.iloc[24][year] + all_rows.iloc[25][year] + all_rows.iloc[26][year] + all_rows.iloc[27][year] + all_rows.iloc[28][year])
        martial_arts =  (all_rows.iloc[13][year] + all_rows.iloc[14][year] + all_rows.iloc[29][year] + all_rows.iloc[40][year] + all_rows.iloc[41][year] + all_rows.iloc[59][year] + all_rows.iloc[67][year] + all_rows.iloc[68][year])
        extreme_sports =(all_rows.iloc[19][year] + all_rows.iloc[20][year] + all_rows.iloc[21][year] + all_rows.iloc[54][year] + all_rows.iloc[55][year] + all_rows.iloc[57][year])
        other_sports =  (all_rows.iloc[5][year] + all_rows.iloc[6][year] + all_rows.iloc[15][year] + all_rows.iloc[16][year] + all_rows.iloc[34][year] + all_rows.iloc[35][year] + all_rows.iloc[36][year] + all_rows.iloc[44][year] + all_rows.iloc[53][year] + all_rows.iloc[61][year] + all_rows.iloc[62][year] + all_rows.iloc[63][year] + all_rows.iloc[64][year] + all_rows.iloc[66][year] + all_rows.iloc[67][year] + all_rows.iloc[68][year] + all_rows.iloc[22][year] + all_rows.iloc[23][year])
        events_sum=aquatics +ball_sports +equestrian + martial_arts +extreme_sports +other_sports

        weighted_sum = total_events
        aquatics_sum = (aquatics/ events_sum).round(3)
        ball_sports_sum = (ball_sports / events_sum).round(3)
        equestrian_sum = (equestrian / events_sum).round(3)
        martial_arts_sum = (martial_arts / events_sum).round(3)
        extreme_sports_sum = (extreme_sports / events_sum).round(3)
        other_sports_sum = (other_sports / events_sum).round(3)

        weighted_sums.append(weighted_sum)
        aquatics_sums.append(aquatics_sum)
        ball_sports_sums.append(ball_sports_sum)
        equestrian_sums.append(equestrian_sum)
        martial_arts_sums.append(martial_arts_sum)
        extreme_sports_sums.append(extreme_sports_sum)
        other_sports_sums.append(other_sports_sum)
    # programs_summary 输出
    programs_summary = pd.DataFrame({
        'Year': year_columns,
        'Total_Programs': weighted_sums,
        'Aquatics_Ratio': aquatics_sums,
        'Ball_Sports_Ratio': ball_sports_sums,
        'Equestrian_Ratio': equestrian_sums,
        'Martial_Arts_Ratio': martial_arts_sums,
        'Extreme_Sports_Ratio': extreme_sports_sums,
    'Other_Sports_Ratio': other_sports_sums
    })
    # 将年份列转换为整数类型
    programs_summary['Year'] = programs_summary['Year'].astype(int)
    # 添加到medal_counts中
    medal_counts = medal_counts.merge(programs_summary, on="Year", how="left")
    # # 输出结果,如果需要保存到文件
    # print(programs_summary)
    # programs_summary.to_csv("programs_summary.csv", index=False, encoding="utf-8")
    # print("处理后的数据已保存到 programs_summary.csv 文件中。")


    ##################################### 运动员获奖比例 ##########################################
    # 是否获奖
    athletes['Medal'] = athletes['Medal'].apply(lambda x: True if x in ['Gold', 'Silver', 'Bronze'] else False)
    athlete_summary = athletes.groupby(["Year", "Team"])["Medal"].mean().reset_index(name="Athlete_Win_Ratio")
    athlete_summary['Athlete_Win_Ratio'] = athlete_summary['Athlete_Win_Ratio'].round(2)
    medal_counts = medal_counts.merge(athlete_summary.rename(columns={"Team": "Country"}), on=["Year", "Country"], how="left").fillna(0)
    # # 输出结果,如果需要保存到文件
    # print(athlete_summary)
    # athlete_summary.to_csv("athlete_summary.csv", index=False, encoding="utf-8")
    # print("处理后的数据已保存到 athlete_summary.csv 文件中。")

    features = medal_counts

    return features



if __name__ == "__main__":
    # 设置文件路径
    file_path = "./data/"  # 替换为你的文件路径

    # 加载数据
    dataframes = input.load_data(file_path)

    # 提取特征
    features = preprocess_data(dataframes)

    # 保存特征至“feature.csv”中
    features.to_csv("./feature.csv", index=False, encoding="utf-8")
    print("处理后的数据已保存到 feature.csv 文件中。")